package simpledb.execution;

import simpledb.common.DbException;
import simpledb.storage.Tuple;
import simpledb.storage.TupleDesc;
import simpledb.transaction.TransactionAbortedException;

import java.util.*;

/**
 * HashEquiJoin 实现了关系数据库中的等值连接（Equi-Join）操作。
 * 它使用哈希表来优化连接过程，避免嵌套循环带来的性能问题。
 */
public class HashEquiJoin extends Operator {

    private static final long   serialVersionUID = 1L;

    // 连接条件谓词
    private final JoinPredicate pred;

    // 两个子迭代器：左表和右表的数据源
    private OpIterator          child1, child2;

    // 合并后的元组描述信息（即连接结果的元组格式）
    private final TupleDesc     comboTD;

    // 缓存当前处理的元组
    transient private Tuple     t1               = null;
    transient private Tuple     t2               = null;

    /**
     * 构造一个等值连接操作符
     * 初始化连接条件和两个数据源
     * 合并两个表的元组描述为一个新的 TupleDesc，用于表示连接结果的结构
     *
     * @param p 连接条件谓词
     * @param child1 左表（外层表）的迭代器
     * @param child2 右表（内层表）的迭代器
     */
    public HashEquiJoin(JoinPredicate p, OpIterator child1, OpIterator child2) {
        this.pred = p;
        this.child1 = child1;
        this.child2 = child2;
        comboTD = TupleDesc.merge(child1.getTupleDesc(), child2.getTupleDesc());
    }

    // 获取器方法。提供对连接条件、输出结构、连接字段名的访问接口

    public JoinPredicate getJoinPredicate() {
        return pred;
    }

    public TupleDesc getTupleDesc() {
        return comboTD;
    }

    public String getJoinField1Name() {
        return this.child1.getTupleDesc().getFieldName(this.pred.getField1());
    }

    public String getJoinField2Name() {
        return this.child2.getTupleDesc().getFieldName(this.pred.getField2());
    }

    // 使用 HashMap 来缓存左表中按连接键分组的元组列表
    final Map<Object, List<Tuple>> map      = new HashMap<>();

    // 设置最大加载记录数，防止内存溢出
    public final static int        MAP_SIZE = 20000;

    /**
     * 将左表（child1）的部分或全部数据加载到哈希表中
     *
     * @return 如果成功加载了 MAP_SIZE 条记录则返回 true，否则返回 false
     * @throws DbException
     * @throws TransactionAbortedException
     */
    private boolean loadMap() throws DbException, TransactionAbortedException {
        int cnt = 0;
        map.clear(); // 清空旧数据

        while (child1.hasNext()) {
            t1 = child1.next();
            Object key = t1.getField(pred.getField1()); // 获取连接键
            List<Tuple> list = map.computeIfAbsent(key, k -> new ArrayList<>());
            list.add(t1); // 添加到对应键的列表中
            if (cnt++ == MAP_SIZE)
                return true; // 达到上限提前返回
        }
        return cnt > 0; // 是否有数据加载
    }

    public void open() throws DbException, NoSuchElementException, TransactionAbortedException {
        child1.open();
        child2.open();
        loadMap(); // 加载左表数据
        super.open();
    }

    public void close() {
        super.close();
        child2.close();
        child1.close();
        this.t1 = null;
        this.t2 = null;
        this.listIt = null;
        this.map.clear(); // 清理资源
    }

    public void rewind() throws DbException, TransactionAbortedException {
        child1.rewind();
        child2.rewind();
    }

    transient Iterator<Tuple> listIt = null;

    /**
     * 处理当前匹配的元组列表，生成连接后的元组
     *
     * @return 合并后的连接元组
     */
    private Tuple processList() {
        t1 = listIt.next();

        int td1n = t1.getTupleDesc().numFields();
        int td2n = t2.getTupleDesc().numFields();

        Tuple t = new Tuple(comboTD);
        for (int i = 0; i < td1n; i++)
            t.setField(i, t1.getField(i));
        for (int i = 0; i < td2n; i++)
            t.setField(td1n + i, t2.getField(i));

        return t;
    }

    /**
     * 获取下一个满足连接条件的元组。
     * 逻辑上是 r1 和 r2 的笛卡尔积中满足连接条件的元组。
     *
     * 此实现采用基于哈希的等值连接算法：
     * - 先构建左表的哈希表；
     * - 然后遍历右表，查找匹配项；
     * - 支持分批加载左表数据。
     *
     * @return 下一个匹配的连接元组，若无更多元组则返回 null
     * @see JoinPredicate#filter
     */
    protected Tuple fetchNext() throws TransactionAbortedException, DbException {
        if (listIt != null && listIt.hasNext()) {
            return processList(); // 返回当前匹配列表中的下一个元组
        }

        // 遍历右表
        while (child2.hasNext()) {
            t2 = child2.next();
            Object key = t2.getField(pred.getField2());

            // 查找哈希表中是否有匹配的左表元组
            List<Tuple> l = map.get(key);
            if (l == null)
                continue;

            listIt = l.iterator(); // 获取匹配列表的迭代器
            return processList(); // 返回第一个匹配元组
        }

        // 右表已处理完，重置右表并尝试加载下一批左表数据
        child2.rewind();
        if (loadMap()) { // 是否还能加载更多左表数据
            return fetchNext(); // 继续处理
        }

        return null; // 无更多数据
    }

    @Override
    public OpIterator[] getChildren() {
        return new OpIterator[] { this.child1, this.child2 };
    }

    @Override
    public void setChildren(OpIterator[] children) {
        this.child1 = children[0];
        this.child2 = children[1];
    }
}
